"""
统一训练脚本
一键训练所有baseline模型
"""
import os
import sys
import argparse
import importlib
import json
import numpy as np
from configs.kfold_config import ALL_MODELS, MODEL_CONFIGS
from train_template import train_model


def main():
    parser = argparse.ArgumentParser(description='训练所有baseline模型')
    parser.add_argument('--models', type=str, nargs='+', default=None,
                        help='指定要训练的模型（默认训练所有）')
    parser.add_argument('--fold', type=int, default=None,
                        help='指定训练哪一折（0-7），不指定则训练所有折')
    parser.add_argument('--resume', action='store_true',
                        help='是否从断点继续训练')
    parser.add_argument('--epochs', type=int, default=None,
                        help='训练轮数')
    parser.add_argument('--batch_size', type=int, default=None,
                        help='批大小')
    parser.add_argument('--lr', type=float, default=None,
                        help='学习率')

    args = parser.parse_args()

    # 确定要训练的模型
    models_to_train = args.models if args.models is not None else ALL_MODELS

    print(f'将训练以下模型: {models_to_train}')
    print(f'{"=" * 80}\n')

    # 存储所有结果
    all_model_results = {}

    for model_name in models_to_train:
        if model_name not in MODEL_CONFIGS:
            print(f'[WARNING] 未知模型: {model_name}，跳过')
            continue

        print(f'\n{"=" * 80}')
        print(f'开始训练模型: {model_name.upper()}')
        print(f'{"=" * 80}\n')

        # 动态导入模型和配置
        model_config = MODEL_CONFIGS[model_name]
        module_name = model_config['module']
        class_name = model_config['class']

        # 导入模型类
        model_module = importlib.import_module(f'{module_name}.model')
        model_class = getattr(model_module, class_name)

        # 导入配置
        config_module = importlib.import_module(f'{module_name}.config')

        # 训练所有fold或指定fold
        if args.fold is not None:
            # 训练指定fold
            print(f'训练 {model_name} - Fold {args.fold}')
            best_metrics = train_model(model_name, model_class, config_module, args.fold, args)
            all_model_results[model_name] = {f'fold_{args.fold}': best_metrics}
        else:
            # 训练所有fold
            k_folds = config_module.TRAIN_CONFIG['k_folds']
            fold_results = {}

            for fold_idx in range(k_folds):
                print(f'\n{"-" * 80}')
                print(f'训练 {model_name} - Fold {fold_idx}')
                print(f'{"-" * 80}\n')

                best_metrics = train_model(model_name, model_class, config_module, fold_idx, args)
                fold_results[f'fold_{fold_idx}'] = best_metrics

            # 计算平均值和标准差
            metrics_names = list(fold_results['fold_0'].keys())
            summary = {}

            for metric_name in metrics_names:
                values = [fold_results[f'fold_{i}'][metric_name] for i in range(k_folds)]
                summary[metric_name] = {
                    'mean': np.mean(values),
                    'std': np.std(values),
                    'values': values
                }

            all_model_results[model_name] = {
                'fold_results': fold_results,
                'summary': summary
            }

            # 打印汇总
            print(f'\n{"=" * 80}')
            print(f'{model_name.upper()} 训练完成！汇总结果：')
            print(f'{"=" * 80}\n')

            for metric_name, stats in summary.items():
                print(f'{metric_name}:')
                print(f'  Mean: {stats["mean"]:.4f}')
                print(f'  Std:  {stats["std"]:.4f}')
                print()

    # 保存所有模型的汇总结果
    summary_path = os.path.join('logs', 'all_models_summary.json')
    os.makedirs(os.path.dirname(summary_path), exist_ok=True)
    with open(summary_path, 'w', encoding='utf-8') as f:
        json.dump(all_model_results, f, indent=4, ensure_ascii=False)

    print(f'\n{"=" * 80}')
    print('所有模型训练完成！')
    print(f'汇总结果已保存到: {summary_path}')
    print(f'{"=" * 80}\n')

    # 打印对比表格
    if args.fold is None:
        print('\n模型性能对比（Mean ± Std）：')
        print(f'{"=" * 80}')
        print(f'{"Model":<15} {"mIoU":<20} {"Dice":<20} {"IoU":<20}')
        print(f'{"-" * 80}')

        for model_name in models_to_train:
            if model_name in all_model_results and 'summary' in all_model_results[model_name]:
                summary = all_model_results[model_name]['summary']
                miou = summary['miou']
                dice = summary['dice']
                iou = summary['iou']

                print(f'{model_name:<15} '
                      f'{miou["mean"]:.4f}±{miou["std"]:.4f}    '
                      f'{dice["mean"]:.4f}±{dice["std"]:.4f}    '
                      f'{iou["mean"]:.4f}±{iou["std"]:.4f}')

        print(f'{"=" * 80}\n')


if __name__ == '__main__':
    main()
